package index

import (
	"gitee.com/Cauchy_AQ/gokv/data"
	"go.etcd.io/bbolt"
	"path/filepath"
)

// BPlusTreeFileName B+ 树索引文件名称
const BPlusTreeFileName = "bptree-index" // TODO: 大写因为 merge 时需要拿到该文件名称，待完成

// indexBucketName B+ 树 bucket 名称
var indexBucketName = []byte("bptree-bucket")

// BPlusTree B+ 树索引
type BPlusTree struct {
	tree *bbolt.DB
	// 无需自定义锁，bbolt.DB 内部已经支持并发
}

// NewBPlusTree 初始化 B+ 树索引
func NewBPlusTree(dirPath string, syncWrites bool) *BPlusTree {
	option := bbolt.DefaultOptions
	option.NoSync = !syncWrites
	bptree, err := bbolt.Open(filepath.Join(dirPath, BPlusTreeFileName), 0644, option)
	if err != nil {
		panic(err)
	}

	if err := bptree.Update(func(tx *bbolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(indexBucketName)
		return err
	}); err != nil {
		panic(err)
	}

	return &BPlusTree{tree: bptree}
}

// Put 向索引中存储 key 对应的数据位置信息
func (b *BPlusTree) Put(key []byte, pos *data.LogRecordPos) *data.LogRecordPos {
	var oldValue []byte
	if err := b.tree.Update(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket(indexBucketName)
		oldValue = bucket.Get(key)
		return bucket.Put(key, data.EncodeLogRecordPos(pos))
	}); err != nil {
		panic(err)
	}
	if len(oldValue) == 0 {
		return nil
	}
	return data.DecodeLogRecordPos(oldValue)
}

// Get 根据 key 取出对于索引位置信息
func (b *BPlusTree) Get(key []byte) *data.LogRecordPos {
	var pos *data.LogRecordPos
	if err := b.tree.View(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket(indexBucketName)
		value := bucket.Get(key)
		if len(value) != 0 {
			pos = data.DecodeLogRecordPos(value)
		}
		return nil
	}); err != nil {
		panic(err)
	}
	return pos
}

// Delete 根据 key 删除对于索引位置信息
func (b *BPlusTree) Delete(key []byte) (*data.LogRecordPos, bool) {
	var oldValue []byte
	if err := b.tree.Update(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket(indexBucketName)
		oldValue = bucket.Get(key)
		if len(oldValue) != 0 {
			return bucket.Delete(key)
		}
		return nil
	}); err != nil {
		panic(err)
	}
	if len(oldValue) == 0 {
		return nil, false
	}
	return data.DecodeLogRecordPos(oldValue), true
}

// Size 获取数据数量大小
func (b *BPlusTree) Size() int {
	var size int
	if err := b.tree.View(func(tx *bbolt.Tx) error {
		bucket := tx.Bucket(indexBucketName)
		size = bucket.Stats().KeyN
		return nil
	}); err != nil {
		panic(err)
	}
	return size
}

// Close 关闭索引
func (b *BPlusTree) Close() error {
	return b.tree.Close()
}

// Iterator 索引迭代器
func (b *BPlusTree) Iterator(reverse bool) Iterator {
	return newBptreeIterator(b.tree, reverse)
}

// bptreeIterator  B+ 树迭代器
type bptreeIterator struct {
	tx           *bbolt.Tx
	cursor       *bbolt.Cursor
	currentKey   []byte
	currentValue []byte
	reverse      bool
}

// newBptreeIterator 初始化 B+ 树迭代器
func newBptreeIterator(tree *bbolt.DB, reverse bool) *bptreeIterator {
	tx, err := tree.Begin(false)
	if err != nil {
		panic(err)
	}
	bpi := &bptreeIterator{
		tx:      tx,
		cursor:  tx.Bucket(indexBucketName).Cursor(),
		reverse: reverse,
	}
	bpi.Rewind()
	return bpi
}

// Rewind 重新返回迭代起点
func (b *bptreeIterator) Rewind() {
	if b.reverse {
		b.currentKey, b.currentValue = b.cursor.Last()
	} else {
		b.currentKey, b.currentValue = b.cursor.First()
	}
}

// Seek 根据传入 key 找到第一个大于（或小于）等于的目标 key，根据这个 key 开始遍历
func (b *bptreeIterator) Seek(key []byte) {
	b.currentKey, b.currentValue = b.cursor.Seek(key)
}

// Next 跳转到下一个 key
func (b *bptreeIterator) Next() {
	if b.reverse {
		b.currentKey, b.currentValue = b.cursor.Prev()
	} else {
		b.currentKey, b.currentValue = b.cursor.Next()
	}
}

// Valid 是否有效，判断是否遍历完了 key
func (b *bptreeIterator) Valid() bool {
	return len(b.currentKey) != 0
}

// Key 获取迭代器当前位置的 key
func (b *bptreeIterator) Key() []byte {
	return b.currentKey
}

// Value 获取当前迭代位置的索引位置信息
func (b *bptreeIterator) Value() *data.LogRecordPos {
	return data.DecodeLogRecordPos(b.currentValue)
}

// Close 关闭迭代器
func (b *bptreeIterator) Close() {
	// bbolt.DB.Begin 为 True，则 Commit；为 False，则 Rollback
	_ = b.tx.Rollback()
}
